package com.zbkj.admin.boot.redisson.ratelimiter;

import com.zbkj.admin.boot.redisson.common.RModule;
import com.zbkj.admin.boot.redisson.common.RedisNameResolver;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RScript;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * redis 限流服务
 *
 * @author dream.lu
 */
@Slf4j
@RequiredArgsConstructor
public class RedisRateLimiterClient implements RateLimiterClient {
	/**
	 * 失败的默认返回值
	 */
	private static final long FAIL_CODE = 0;
	/**
	 * RedissonClient
	 */
	private final RScript script;
	/**
	 * luaScript
	 */
	private final String luaScript;
	private final RedisNameResolver resolver;

	@Override
	public boolean isAllowed(String key, long max, long ttl, TimeUnit timeUnit) {
		List<Object> keys = getKeys(key);
		// 转为毫秒
		long ttlSeconds = timeUnit.toMillis(ttl);
		// 令牌桶总容量
		Long results = script.eval(RScript.Mode.READ_WRITE, luaScript, RScript.ReturnType.MULTI, keys, max, ttlSeconds, 1);
		// 结果为空返回失败
		if (results == null) {
			return false;
		}
		// 判断返回成功
        return results != FAIL_CODE;
	}

	private List<Object> getKeys(String id) {
		// redis key
		String redisPrefix = resolver.resolve(RModule.RateLimiter, id);
		// use `{}` around keys to use Redis Key hash tags
		// this allows for using redis cluster
		// Make a unique key per user.
		String prefix = redisPrefix + ".{" + id;
		// You need two Redis keys for Token Bucket.
		String tokenKey = prefix + "}.tokens";
		String timestampKey = prefix + "}.timestamp";
		return Arrays.asList(tokenKey, timestampKey);
	}

}
